import os
os.environ["PROJ_LIB"] = r'C:\Users\ckpark\anaconda3\Library\share'

import numpy as np
import matplotlib.pyplot as plt
import numpy.ma as ma
from mpl_toolkits.basemap import Basemap
from scipy.stats import linregress


def init_plotting():
    plt.rcParams['font.family'] = 'Times New Roman'
init_plotting()

read = np.loadtxt('water_rev_2.txt', dtype='str')
lon = np.array(read[:,0], dtype='float')
lat = np.array(read[:,1], dtype='float')
dat = np.array(read[:,2], dtype='float')
sig = np.array(read[:,3], dtype='int')
stor = np.array(read[:,4], dtype='float')

mstor = np.mean(stor)

smin = np.min(stor)
smax = np.max(stor)

wstor = np.zeros((len(stor)))
for t in range(len(stor)):
    wstor[t] = np.log10(stor[t])**4

print(wstor)

slat = 34.2; elat = 35.6
slon = 126; elon = 127.8

fig_name = 'spatial_water_rev_map.png'
#title=['Cases','Mean intensity','Mean duration','Proportion of'+'\n'+'rapid developing drought']
fig = plt.figure(figsize=(15,5))
for xx in range(1):
    ax = fig.add_subplot(1,1,xx+1)
    
    s = wstor.argsort()
    odat = dat[s[::-1]]
    wdat = wstor[s[::-1]]
    olat = lat[s[::-1]]
    olon = lon[s[::-1]]
 
    cmax = 100
    cmin = 30
    cint = 8
    jet = plt.cm.get_cmap('jet_r')

    levels = np.linspace(cmin,cmax,cint)

    m = Basemap(projection='mill',llcrnrlat=slat,urcrnrlat=elat,llcrnrlon=slon,urcrnrlon=elon,lat_ts=20,resolution='i')
    m.drawcoastlines(color='black')
    parallels = np.arange(33,elat,0.5)
    m.drawparallels(parallels,labels=[True,False,False,False],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    meridians = np.arange(126,elon,1) 
    m.drawmeridians(meridians,labels=[True,False,False,True],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    m.drawcountries(color='black') 
    #m.drawstates(color='black')

    x,y = m(olon, olat)

    #plt.title(title[xx],fontsize=20,loc='center')

    for t in range(len(olat)):
        ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet, edgecolor='none', s=int(2*wdat[t]),zorder=1,alpha=1.0)

    cb = fig.colorbar(ax11,ticks=levels,fraction=0.035,pad=0.025,extend='both')
    cb.ax.tick_params(labelsize=15)

plt.subplots_adjust(hspace=0.3,wspace=0.35)
plt.show() 
fig.savefig(fig_name,bbox_inches='tight')
####################
slat = 34.2; elat = 35.6
slon = 126; elon = 127.8


lstor = [1000, 10000, 100000]
ldat = [1, 1, 1]
wlstor = np.log10(lstor)**4
print(wlstor)

llat = lat[:9][::3]
llon = lon[:9][::3]

print(10**8)

fig_name = 'legend_water_rev_map.png'
#title=['Cases','Mean intensity','Mean duration','Proportion of'+'\n'+'rapid developing drought']
fig = plt.figure(figsize=(15,5))
for xx in range(1):
    ax = fig.add_subplot(1,1,xx+1)
    
    odat = ldat
    wdat = wlstor
    olat = llat
    olon = llon
 
    cmax = 1
    cmin = 0
    cint = 8
    jet = plt.cm.get_cmap('binary')

    levels = np.linspace(cmin,cmax,cint)

    m = Basemap(projection='mill',llcrnrlat=slat,urcrnrlat=elat,llcrnrlon=slon,urcrnrlon=elon,lat_ts=20,resolution='i')
    m.drawcoastlines(color='black')
    parallels = np.arange(33,elat,0.5)
    m.drawparallels(parallels,labels=[True,False,False,False],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    meridians = np.arange(126,elon,1) 
    m.drawmeridians(meridians,labels=[True,False,False,True],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    m.drawcountries(color='black') 
    #m.drawstates(color='black')

    x,y = m(olon, olat)

    #plt.title(title[xx],fontsize=20,loc='center')
    for t in range(len(olat)):
        #if sig[t] == 1:
        #    ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet,marker='*', edgecolor='none', s=600,zorder=10)
        #    ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet,marker='*', facecolor='none',edgecolor='black', s=600,zorder=10)
        #else:
        ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet, edgecolor='none', s=int(wdat[t]*2),zorder=1,alpha=1.0)

    cb = fig.colorbar(ax11,ticks=levels,fraction=0.035,pad=0.025,extend='both')
    cb.ax.tick_params(labelsize=15)

plt.subplots_adjust(hspace=0.3,wspace=0.35)
plt.show() 
fig.savefig(fig_name,bbox_inches='tight')
#################



slat = 33; elat = 44
slon = 124; elon = 131

'''
fig_name = '../99.FIG/dummy_kor.png'
#title=['Cases','Mean intensity','Mean duration','Proportion of'+'\n'+'rapid developing drought']
fig = plt.figure(figsize=(15,5))
for xx in range(1):
    ax = fig.add_subplot(1,1,xx+1)
    odat = dat
	
    cmax = 100
    cmin = 30
    cint = 8
    jet = plt.cm.get_cmap('jet_r')

    levels = np.linspace(cmin,cmax,cint)

    m = Basemap(projection='mill',llcrnrlat=slat,urcrnrlat=elat,llcrnrlon=slon,urcrnrlon=elon,lat_ts=20,resolution='i')
    m.drawcoastlines(color='black')
    parallels = np.arange(33,elat,2)
    m.drawparallels(parallels,labels=[True,False,False,False],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    meridians = np.arange(124,elon,3) 
    m.drawmeridians(meridians,labels=[True,False,False,True],linewidth=0.1,fontsize=22)
    m.drawmapboundary(fill_color='white')
    m.drawcountries(color='black') 
    #m.drawstates(color='black')

    x,y = m(lon, lat)

    #plt.title(title[xx],fontsize=20,loc='center')
    #for t in range(len(lat)):
    #    if sig[t] == 1:
    #        ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet,marker='*', edgecolor='none', s=600,zorder=10)
    #        ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet,marker='*', facecolor='none',edgecolor='black', s=600,zorder=10)
    #    else:
    #        ax11 = m.scatter(x[t],y[t],c=odat[t], vmin=cmin, vmax=cmax, cmap=jet, edgecolor='none', s=220,zorder=1)

    #cb = fig.colorbar(ax11,ticks=levels,fraction=0.035,pad=0.025,extend='both')
    #cb.ax.tick_params(labelsize=15)

plt.subplots_adjust(hspace=0.3,wspace=0.35)
plt.show() 
fig.savefig(fig_name,bbox_inches='tight')
'''